package com.github.kezhenxu94.data.sink;

import org.apache.solr.client.solrj.SolrClient;
import org.apache.solr.client.solrj.SolrServerException;
import org.apache.solr.client.solrj.impl.HttpSolrClient;
import org.apache.solr.common.SolrInputDocument;

import java.io.IOException;
import java.util.Properties;
import java.util.concurrent.*;
import java.util.function.Function;

/**
 * Created by kezhenxu on 1/7/17.
 *
 * @author kezhenxu (kezhenxu94@163.com)
 */
public class SolrSink<T> implements Sink<T> {

    protected ExecutorService executor;
    protected Properties properties;
    protected BlockingQueue<T> queue;
    protected volatile boolean isRunning;
    protected SolrClient client;
    protected String collection;
    protected Function<T, SolrInputDocument> valueSupplier;

    @Override
    public void setup() throws Exception {
        this.executor = Executors.newSingleThreadExecutor();
        this.properties = new Properties();
        this.properties.load(getClass().getResourceAsStream("/solr-sink.properties"));
        this.client = new HttpSolrClient.Builder(this.properties.getProperty("solr.home")).build();
        this.queue = new LinkedBlockingQueue<>();
    }

    @Override
    public void start() throws Exception {
        isRunning = true;
        executor.execute(() -> {
            while (isRunning()) {
                try {
                    T take = queue.take();
                    client.add(collection, valueSupplier.apply(take));
                } catch (InterruptedException | SolrServerException | IOException e) {
                    getLogger().error(e.getMessage(), e);
                }
            }
        });
    }

    @Override
    public void close() throws Exception {
        try {
            isRunning = false;
            while (isRunning()) {
                TimeUnit.MILLISECONDS.sleep(500L);
            }
        } finally {
            client.close();
        }
    }

    @Override
    public boolean isRunning() {
        return isRunning || !queue.isEmpty();
    }

    @Override
    public void put(T t) throws Exception {
        if (!isRunning) {
            throw new IllegalStateException("Sink has been closed. ");
        }
        queue.offer(t);
    }

    public String getCollection() {
        return collection;
    }

    public void setCollection(String collection) {
        this.collection = collection;
    }

    public Function<T, SolrInputDocument> getValueSupplier() {
        return valueSupplier;
    }

    public void setValueSupplier(Function<T, SolrInputDocument> valueSupplier) {
        this.valueSupplier = valueSupplier;
    }
}
